import numpy as np
import onnxruntime as ort
from onnx import numpy_helper
import onnx

def test_conv(conv_in:str, conv_in_size:list, conv_weight:str, weight_size:list, conv_bias:str, strides:list, pads:list, dilations:list, group:int=1):
    assert len(conv_in_size) == 4 and len(weight_size) == 4, "only support input dimension is 4!"
    in_arr = np.fromfile(conv_in, dtype=np.float32).reshape(conv_in_size)
    weight_arr = np.fromfile(conv_weight, dtype=np.float32).reshape(weight_size)
    weight_tensor = numpy_helper.from_array(weight_arr, "w")
    input_names = ["x", "w"]
    init = [weight_tensor]
    if conv_bias:
        bias_arr = np.fromfile(conv_bias, dtype=np.float32)
        bias_tensor = numpy_helper.from_array(bias_arr, "b")
        init.append(bias_tensor)
        input_names.append("b")
    input_x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, conv_in_size)
    output_y = onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, None)

    conv_node = onnx.helper.make_node(
        "Conv",
        inputs=input_names,
        outputs=["y"],
        kernel_shape=weight_size[-2:],
        strides = strides,
        pads=[1, 1, 1, 1],
        dilations=dilations,
        group=group
    )

    graph = onnx.helper.make_graph(
        nodes=[conv_node], name="onnx_sample", 
        inputs=[input_x], outputs=[output_y],
        initializer=init,
    )

    model = onnx.helper.make_model(graph, producer_name="alibaba")
    model = onnx.shape_inference.infer_shapes(model)
    onnx.save(model, "./data/conv.onnx")

    # Checking the produces are the expected ones.
    sess = ort.InferenceSession(model.SerializeToString(),
                                        providers=["CPUExecutionProvider"])
    feeds = {"x": in_arr}
    results = sess.run(["y"], feeds)[0]
    results.tofile("./data/conv_output.bin")
    print(f"{type(results)}, output shape = {results.shape}")


    print("done")



if __name__ == "__main__":
    conv_in_size = [1,32,224,224]
    conv_w_size = [32, 2, 3, 3]
    """
    conv_in = np.random.randint(-64, 64, conv_in_size, dtype=np.int8).astype(np.float32)
    conv_w = np.random.random(conv_w_size).astype(np.float32)
    conv_b = np.random.random([conv_w_size[0]]).astype(np.float32)
    conv_in.tofile("./data/conv_in.bin")
    conv_w.tofile("./data/conv_w.bin")
    conv_b.tofile("./data/conv_b.bin")
    """

    test_conv("./data/conv_in.bin", conv_in_size, 
               "./data/conv_w.bin", conv_w_size,
                "./data/conv_b.bin",
                # None,
                [1,1],
                [1,1,1,1],
                [1,1],
                16)